#%% config the app
# exp_no = '2024-1-18-1'
exp_no = '01'

# mode = 'train'
mode = 'test'


# folder = "C:\\Users\\86198\\Desktop\\8针水表\\8针水表\\8P Uneven"train_segmentor.py
# folder = "./img/8P Uneven"
folder = "./img/8P Even"




#%% import libs
# import libs
import os
import sys
import numpy as np
import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from torchvision.models.detection import MaskRCNN
import torchvision.transforms as tT
import torchvision.transforms.functional as F
from torchvision.utils import draw_segmentation_masks
import matplotlib.pyplot as plt

from PIL import Image, ExifTags
import json
import PySimpleGUI as sg
from _pointer_meter_helpers import rotate_im_accord_exiftag, load_anno, load_valid_imfile_names

sys.path.insert(0, './torchvision_det_references') #确保可以通过下面的语句导入位于子目录中的包
import utils
from engine import train_one_epoch, evaluate
import transforms as T

# %% 定义dataset类
class MeterSegDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms, down_scale_factor=8):#down_scale_factor比例缩减因子
        self.root = root # root folder path
        self.transforms = transforms # data transformations
        self.down_scale_factor = down_scale_factor # 图像缩小为原图的比例
                            
        self.imgs = load_valid_imfile_names(root)
    
    def __len__(self):
        return len(self.imgs)
    
    def __getitem__(self, idx):        
        fpath = os.path.join(self.root, self.imgs[idx])
        img = Image.open(fpath)
        img = rotate_im_accord_exiftag(img)#处理手机拍照时的旋转问题
        img = img.convert('RGB')#如果不使用.convert('RGB')进行转换的话，读出来的图像是RGBA四通道的，A通道为透明通道
        #缩小图像
        im_sz=(img.size[0]//self.down_scale_factor, img.size[1]//self.down_scale_factor)#这里size返回的是宽高而不是像素点多少
        img = img.resize(im_sz)
        
        #载入标注信息
        anno = load_anno(self.root, self.imgs[idx], self.down_scale_factor)
        #获得包围盒信息
        boxes = []
        for b in anno['mask_boxes']: #遍历标注信息的mask_boxes（需要可以去原函数把注释解掉）
            boxes.append(b)
        
        boxes = torch.as_tensor(boxes, dtype=torch.float32)

        #获得类别信息，只有1类目标就是仪表
        num_objs = len(boxes)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(anno['masks'], dtype=torch.uint8)
        '''当mask的数据类型是torch.uint8时，此时的tensor用作mask，tensor中的1对应的行/列保留，
        0对应的行/列舍去。且被mask的维度必须与原始tensor的维度一致.
        当masks的数据类型是torch.long时，此时的tensor用作index，tensor中的每个数字代表着将要取出的tensor的行列索引'''
        #获取其他coco格式的必要信息
        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        #target字典获得标签信息
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        #进行变化
        if self.transforms is not None:
            img, target = self.transforms(img, target)
        
        return img, target        

# %% define the mask faster-rcnn -based model
num_classes = 2 #我们数据集共2个类别，背景和指针
# load an instance segmentation model pre-trained on COCO(实例分割模型）
model = torchvision.models.detection.maskrcnn_resnet50_fpn(weights="DEFAULT")

# 获取分类器的输入特征数
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# 现在获取掩码分类器的输入特征数量
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,
                                                       num_classes)

# %% prepare for training

def get_transform(train):
    transforms = []
    transforms.append(T.PILToTensor())
    transforms.append(T.ConvertImageDtype(torch.float))
    if train:
        transforms.append(T.RandomPhotometricDistort())
    return T.Compose(transforms)

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

dataset = MeterSegDataset(folder, get_transform(train=True), down_scale_factor=8)
dataset_test = MeterSegDataset(folder, get_transform(train=False),down_scale_factor=8)

# 划分训练集和测试集的数据
indices = torch.randperm(len(dataset)).tolist()#randperm（）随机打乱函数
dataset = torch.utils.data.Subset(dataset, indices[:-30])# Subset获取指定一个索引序列对应的子数据集
dataset_test = torch.utils.data.Subset(dataset_test, indices[-30:])



# import torch.utils.data#不加这个下面data报错！！！！
'''DataLoader参数：
dataset (Dataset) – 加载数据的数据集。
batch_size (int, optional) – 每个batch加载多少个样本(默认: 1)。
shuffle (bool, optional) – 设置为True时会在每个epoch重新打乱数据(默认: False).
sampler (Sampler, optional) – 定义从数据集中提取样本的策略，即生成index的方式，可以顺序也可以乱序
num_workers (int, optional) – 用多少个子进程加载数据。0表示数据将在主进程中加载(默认: 0)
collate_fn (callable, optional) –将一个batch的数据和标签进行合并操作。
pin_memory (bool, optional) –设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
drop_last (bool, optional) – 如果数据集大小不能被batch size整除，则设置为True后可删除最后一个不完整的batch。如果设为False并且数据集的大小不能被batch size整除，则最后一个batch将更小。(默认: False)
timeout，是用来设置数据读取的超时时间的，但超过这个时间还没读取到数据的话就会报错'''
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=2, shuffle=True, num_workers=0,
    collate_fn=utils.collate_fn)

data_loader_test = torch.utils.data.DataLoader(
    dataset_test, batch_size=1, shuffle=False, num_workers=0,
    collate_fn=utils.collate_fn)

# move model to the right device
model.to(device)

# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
                            momentum=0.9, weight_decay=0.0005)

'''学习率调整
orch.optim.lr_scheduler.StepLR(optimizer,step_size,gamma=0.1,last_epoch=-1,verbose=False)
optimizer模型优化器
step_size学习率更新步长(注意这里是每step_size个epoch更新一次)
gamma 衰减率(默认0.1 即每step_size个epoch后lr变为lr*gamma
last_epoch last_epoch之后恢复lr为initial_lr(如果是训练了很多个epoch后中断了 继续训练 这个值就等于加载的模型的epoch 默认为-1表示从头开始训练，即从epoch=1开始。
verbose 是否每次改变都输出一次lr的值(单词的意思就是“冗长的” 默认是False'''
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=3,
                                                gamma=0.1)
# %% do training
fn = './weights/model_weights_seg_'+exp_no+'.pth'

if mode == "train":
    num_epochs = 10

    for epoch in range(num_epochs):
        # train for one epoch, printing every 10 iterations
        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
        # update the learning rate
        lr_scheduler.step()
        # evaluate on the test dataset
        evaluate(model, data_loader_test, device=device)

    print("Done!")

    # %% save the model    
    torch.save(model.state_dict(), fn) 

# %% test the model
if mode == 'test':
    model.load_state_dict(torch.load(fn))
    model.eval()
    '''我们需要将这些层设置到预测模式，model.eval()就是帮我们一键搞定的,
    相应的，在训练之前，我们也要记得将这些特殊的层设置到训练模式：model.train()'''

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    model.to(device)

plt.rcParams["savefig.bbox"] = 'tight'#自动调整图形周围的填充

def show(imgs):
    if not isinstance(imgs, list):#isinstance() 函数来判断一个对象是否是一个已知的类型，此处是判断imgs是否为列表类型
        imgs = [imgs]
    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    '''上面是这些的简写：
    fig = plt.figure()#首先调用plt.figure()创建了一个**画窗对象fig**
ax = fig.add_subplot(111)#然后再对fix创建默认的坐标区（一行一列一个坐标区）
#这里的（111）相当于（1，1，1），当然，官方有规定，当子区域不超过9个的时候，我们可以这样简写
'''
    for i, img in enumerate(imgs):
        img = img.detach()#返回一个新的tensor，从当前计算图中分离下来。但是仍指向原变量的存放位置，不同之处只是requirse_grad为false.得到的这个tensir永远不需要计算器梯度，不具有grad.
        img = F.to_pil_image(img)#转为pil img
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

# 界面布局，将会按照列表顺序从上往下依次排列，二级列表中，从左往右依此排列
layout = [  [sg.Text('共有'+str(len(dataset_test))+"个带标注的水表图像")],
            [sg.Text('请输入要查看的图像序号(从0开始):'), sg.InputText('0')],   #这里默认值设置为‘0’
            [sg.Button('Ok'), sg.Button('Cancel')] ]

# Create the Window
window = sg.Window('检测结果查看程序', layout)
plt.ion()#打卡动态画图 i on
fig = plt.figure()
# Event Loop to process "events" and get the "values" of the inputs
while True:
    event, values = window.read()
    if event == sg.WIN_CLOSED or event == 'Cancel': # if user closes window or clicks cancel
        break
    if(int(values[0])>=0 and int(values[0])<len(dataset_test)):
        fid = int(values[0])
        img,label = dataset_test.__getitem__(fid)
        
        x = img.unsqueeze(0)#.unsqueeze(0)增加维度（0表示，在第一个位置增加维度）
        x = x.to(device)
        predictions = model(x)  
        
        plt.cla()## 清除axes，即当前 figure 中的活动的axes，但其他axes保持不变
        #show the original image
        mask = predictions[0]['masks']
        # pil_img = tT.ToPILImage()(img)
        pil_img = tT.ToPILImage()(mask[0])
        plt.imshow(pil_img)

        #显示包围盒的标签信息
        boxes_gt = label['boxes'].numpy()
        nbox = boxes_gt.shape[0]
        for i in range(nbox):
            box = boxes_gt[i,:]
            w = box[2]-box[0]
            h = box[3]-box[1]
            #gca就是get current axes  ，add_patch()函数用于将补丁添加到轴的补丁；返回补丁
            plt.gca().add_patch(plt.Rectangle((box[0],box[1]), w,h,fill=False, color='green'))


        #显示包围盒        
        boxes = predictions[0]['boxes'].cpu().detach().numpy()
        nbox = boxes.shape[0]
        for i in range(nbox):
            box = boxes[i,:]
            w = box[2]-box[0]
            h = box[3]-box[1]
            plt.gca().add_patch(plt.Rectangle((box[0],box[1]), w,h,fill=False, color='red'))
        
        
        # draw_im_with_anno(folder, flst[fid], anno)
        plt.draw()

window.close()

# %%
